package com.tfswx.fljdmxyy.utils.sm4;

import java.io.ByteArrayInputStream;
import java.io.InputStream;

public class EncryptSm4InputStream extends DecorateInputStreamBase {

    private static final InnerSM4_Context SM4_CONTEXT_ENCRYPT = new InnerSM4_Context(InnerSM4Stream.SM4_ENCRYPT);
    private final byte[] iv = new byte[16];
    private final byte[] out = new byte[16];
    private volatile int inputLength = 0;
    private volatile boolean isEnd;
    public EncryptSm4InputStream(InputStream in) {
        super(in, 16, 16);

        System.arraycopy(SM4_CONTEXT_ENCRYPT.ivBytes, 0, iv, 0, 16);
    }
    public EncryptSm4InputStream(byte[] decryptData) {
        this(createInputStream(decryptData));
    }

    private static InputStream createInputStream(byte[] decryptData) {
        if (decryptData == null || decryptData.length == 0) {
            throw new IllegalArgumentException("decryptData 不能为空");
        }

        return new ByteArrayInputStream(decryptData);
    }

    @Override
    protected int innerReadTo(byte[] sourceBuffer, int sourceBufferLength, byte[] destBuffer) {

        if (sourceBufferLength > 0) {
            inputLength += sourceBufferLength;
        }

        if (sourceBufferLength == 16) {
            // 直接加密
            for (int i = 0; i < 16; i++) {
                out[i] = ((byte) (sourceBuffer[i] ^ iv[i]));
            }
            InnerSM4Stream.sm4_one_round(SM4_CONTEXT_ENCRYPT.sk, out, destBuffer);
            System.arraycopy(destBuffer, 0, iv, 0, 16);
            return 16;
        } else if (sourceBufferLength <= 0) {
            // 判断上一把是否已经读完
            if (isEnd) {
                return -1;
            }
            isEnd = true;

            // 上一把刚好16长度，本轮需要补位
            InnerSM4Stream.paddingByteArray16(sourceBuffer, inputLength);

            for (int i = 0; i < 16; i++) {
                out[i] = ((byte) (sourceBuffer[i] ^ iv[i]));
            }
            InnerSM4Stream.sm4_one_round(SM4_CONTEXT_ENCRYPT.sk, out, destBuffer);
            return 16;
        } else {
            // 末尾不足16
            if (isEnd) {
                return -1;
            }
            isEnd = true;
            InnerSM4Stream.paddingByteArray16(sourceBuffer, inputLength);

            for (int i = 0; i < 16; i++) {
                out[i] = ((byte) (sourceBuffer[i] ^ iv[i]));
            }
            InnerSM4Stream.sm4_one_round(SM4_CONTEXT_ENCRYPT.sk, out, destBuffer);
            return 16;
        }
    }
}
